package drds.plus.datanode.api;

import drds.plus.common.jdbc.SqlTypeParser;
import drds.plus.common.model.SqlType;
import drds.plus.datanode.configuration.DataSourceIndexAndNeedRetryIfDailed;
import drds.plus.datanode.configuration.DataSourceWrapper;
import drds.plus.datanode.executor.AbstractExecutor;
import drds.plus.datanode.executor.Executor;
import drds.plus.datasource.api.ConnectionWrapper;
import drds.plus.datasource.api.ResultSetWrapper;
import drds.plus.datasource.api.StatementWrapper;
import lombok.extern.slf4j.Slf4j;

import java.sql.SQLException;
import java.util.LinkedList;
import java.util.List;

@Slf4j
public class Statement {
    private final String applicationId;
    protected DatasourceManager datasourceManager;
    protected Connection connection;
    public StatementWrapper statementWrapper;

    protected int retryingTimes;
    /**
     * query time out . 超时时间，如果超时时间不为0。那么超时应该被set到真正的query中。
     */
    protected int queryTimeout = 0;
    protected int fetchSize;
    protected int maxRows;
    /**
     * 经过计算后的结果集，允许使用 getResult函数调用. 一个statement只允许有一个结果集
     */
    protected ResultSetWrapper resultSetWrapper;
    /**
     * 更新计数，如果执行了多次，那么这个值只会返回最后一次执行的结果。 如果是一个query，那么返回的数据应该是-1
     */
    protected int updateCount;

    protected List<String> batchedArgs;

    protected boolean closed;
    /**
     * 貌似是只有存储过程中会出现多结果集 因此不支持
     */
    protected boolean moreResults;
    protected Executor<ResultSetWrapper> executeQueryExecutor = new AbstractExecutor<ResultSetWrapper>() {

        public ResultSetWrapper createConnectionWrapperAndExecute(DataSourceWrapper dataSourceWrapper, Object... args) throws SQLException {
            String sql = (String) args[0];
            ConnectionWrapper connectionWrapper = Statement.this.connection.createConnectionWrapper(dataSourceWrapper, true);
            return executeQuery(connectionWrapper, sql);
        }
    };
    private Executor<Integer> executeUpdateExecutor = new AbstractExecutor<Integer>() {

        public Integer createConnectionWrapperAndExecute(DataSourceWrapper dataSourceWrapper, Object... args) throws SQLException {
            ConnectionWrapper connectionWrapper = Statement.this.connection.createConnectionWrapper(dataSourceWrapper, false);
            return executeUpdate(connectionWrapper, (String) args[0], (Integer) args[1]);
        }
    };
    private Executor<int[]> executeBatchExecutor = new AbstractExecutor<int[]>() {

        public int[] createConnectionWrapperAndExecute(DataSourceWrapper dataSourceWrapper, Object... args) throws SQLException {
            ConnectionWrapper connectionWrapper = Statement.this.connection.createConnectionWrapper(dataSourceWrapper, false);
            return executeBatch(connectionWrapper, Statement.this.batchedArgs);
        }
    };

    public Statement(String applicationId, DatasourceManager datasourceManager, Connection connection) {
        this.applicationId = applicationId;
        this.datasourceManager = datasourceManager;
        this.connection = connection;
        this.retryingTimes = datasourceManager.getRetryingTimes();
    }

    /**
     * 设置在底层执行的具体的Statement 如果前面的baseStatement未关，则先关闭
     */
    void setStatementWrapper(StatementWrapper statementWrapper) {
        if (this.statementWrapper != null) {
            try {
                this.statementWrapper.close();
            } catch (SQLException e) {
                log.error("close statementWrapper failed.", e);
            }
        }
        this.statementWrapper = statementWrapper;
    }

    public boolean execute(String sql) throws SQLException {
        return executeInternal(sql, -1);
    }

    public boolean execute(String sql, int autoGeneratedKeys) throws SQLException {
        return executeInternal(sql, autoGeneratedKeys);
    }

    // jdbc规范: 返回true表示executeQuery，false表示executeUpdate
    private boolean executeInternal(String sql, int autoGeneratedKeys) throws SQLException {
        if (SqlTypeParser.isQuerySql(sql)) {
            executeQuery(sql);
            return true;
        } else {
            if (autoGeneratedKeys != -1) {
                executeUpdate(sql, autoGeneratedKeys);
            } else {
                executeUpdate(sql);
            }
            return false;
        }
    }

    /*
     * ========================================================================
     * executeUpdate逻辑
     * ======================================================================
     */
    public int executeUpdate(String sql) throws SQLException {
        return executeUpdateInternal(sql, -1);
    }

    public int executeUpdate(String sql, int autoGeneratedKeys) throws SQLException {
        return executeUpdateInternal(sql, autoGeneratedKeys);
    }

    private int executeUpdateInternal(String sql, int autoGeneratedKeys) throws SQLException {
        checkClosed();
        ensureResultSetWrapperIsEmpty();
        ConnectionWrapper connectionWrapper = connection.getConnectionWrapper(sql, false);
        if (connectionWrapper != null) {
            sql = DataNodeHintParser.removeDataNodeIdHint(sql);
            this.updateCount = executeUpdate(connectionWrapper, sql, autoGeneratedKeys);
            return this.updateCount;
        } else {
            DataSourceIndexAndNeedRetryIfDailed dataSourceIndexAndNeedRetryIfDailed = DataNodeHintParser.convertHint2Index(sql);
            sql = DataNodeHintParser.removeDataNodeIdHint(sql);
            if (dataSourceIndexAndNeedRetryIfDailed == null) {
                dataSourceIndexAndNeedRetryIfDailed = ThreadLocalDataSourceIndex.getIndex();
            }
            this.updateCount = this.datasourceManager.getSelector(false).getConnectionWrapperAndExecuteAndRetryIfFailed(null, executeUpdateExecutor, retryingTimes, sql, autoGeneratedKeys, dataSourceIndexAndNeedRetryIfDailed);
            return this.updateCount;
        }
    }

    private int executeUpdate(ConnectionWrapper connectionWrapper, String sql, int autoGeneratedKeys) throws SQLException {
        StatementWrapper statementWrapper = createStatementWrapper(connectionWrapper, sql, false);
        if (autoGeneratedKeys != -1) {
            return statementWrapper.executeUpdate(sql, autoGeneratedKeys);
        } else {
            return statementWrapper.executeUpdate(sql);
        }
    }

    /**
     * 会调用setBaseStatement以关闭已有的Statement
     */
    private StatementWrapper createStatementWrapper(ConnectionWrapper connectionWrapper, String sql, boolean isBatch) throws SQLException {
        StatementWrapper statementWrapper = connectionWrapper.createStatement();
        setStatementWrapper(statementWrapper); // 会关闭已有的Statement
        statementWrapper.setQueryTimeout(queryTimeout); // 这句也有可能抛出异常，放在最后
        statementWrapper.setFetchSize(fetchSize);
        statementWrapper.setMaxRows(maxRows);
        return statementWrapper;
    }

    public void addBatch(String sql) throws SQLException {
        checkClosed();
        if (batchedArgs == null) {
            batchedArgs = new LinkedList<String>();
        }
        if (sql != null) {
            batchedArgs.add(sql);
        }
    }


    public int[] executeBatch() throws SQLException {
        try {
            checkClosed();
            ensureResultSetWrapperIsEmpty();

            if (batchedArgs == null || batchedArgs.isEmpty()) {
                return new int[0];
            }

            ConnectionWrapper connectionWrapper = connection.getConnectionWrapper(null, false);
            if (connectionWrapper != null) {
                // 如果当前已经有连接,则不做任何重试。对于更新来说，不管有没有事务，
                // 用户总期望getConnection获得连接之后，后续的一系列操作都在这同一个库，同一个连接上执行
                return executeBatch(connectionWrapper, this.batchedArgs);
            } else {
                return datasourceManager.getSelector(false).getConnectionWrapperAndExecuteAndRetryIfFailed(null, executeBatchExecutor, retryingTimes);
            }
        } finally {
            if (batchedArgs != null) {
                batchedArgs.clear();
            }
        }
    }

    private int[] executeBatch(ConnectionWrapper connectionWrapper, List<String> batchedSqls) throws SQLException {
        StatementWrapper statementWrapper = createStatementWrapper(connectionWrapper, batchedSqls.get(0), true);
        for (String sql : batchedSqls) {
            statementWrapper.addBatch(sql);
        }
        return statementWrapper.executeBatch();
    }

    public void clearBatch() throws SQLException {
        checkClosed();
        if (batchedArgs != null) {
            batchedArgs.clear();
        }
    }

    //
    public void close() throws SQLException {
        close(true);
    }

    void close(boolean removeThis) throws SQLException {
        if (closed) {
            return;
        }
        closed = true;
        if (resultSetWrapper != null) {
            try {
                resultSetWrapper.close();
            } catch (SQLException e) {
                log.warn("Close resultSetWrapper failed.", e);
            } finally {
                resultSetWrapper = null;
            }
        }
        if (this.statementWrapper != null) {
            try {
                this.statementWrapper.close();
            } catch (SQLException e) {
                log.warn("Close statementWrapper failed.", e);
            }
        }
        this.statementWrapper = null;
        if (removeThis) {
            connection.removeOpenedStatements(this);
        }
    }

    protected void checkClosed() throws SQLException {
        if (closed) {
            throw new SQLException("No operations allowed execute_plan_optimizer statement closed.");
        }
    }

    /**
     * 如果新建了查询，那么上一次查询的结果集应该被显示的关闭掉。这才是符合jdbc规范的
     */
    protected void ensureResultSetWrapperIsEmpty() throws SQLException {
        if (resultSetWrapper != null) {
            try {
                resultSetWrapper.close();
            } catch (SQLException e) {
                log.error("exception on close last result set . can do nothing..", e);
            } finally {
                resultSetWrapper = null;
            }
        }

    }

    /*
     * ========================================================================
     * executeQuery 查询逻辑
     * ======================================================================
     */
    public ResultSetWrapper executeQuery(String sql) throws SQLException {
        checkClosed();
        ensureResultSetWrapperIsEmpty();
        boolean read = SqlType.SELECT.equals(SqlTypeParser.getSqlType(sql)) && connection.getAutoCommit();
        ConnectionWrapper connectionWrapper = connection.getConnectionWrapper(sql, read);
        if (connectionWrapper != null) {
            sql = DataNodeHintParser.removeDataNodeIdHint(sql);
            return executeQuery(connectionWrapper, sql);
        } else {
            // hint优先
            DataSourceIndexAndNeedRetryIfDailed dataSourceIndexAndNeedRetryIfDailed = DataNodeHintParser.convertHint2Index(sql);
            sql = DataNodeHintParser.removeDataNodeIdHint(sql);
            if (dataSourceIndexAndNeedRetryIfDailed == null) {
                dataSourceIndexAndNeedRetryIfDailed = ThreadLocalDataSourceIndex.getIndex();
            }
            return this.datasourceManager.getSelector(read).getConnectionWrapperAndExecuteAndRetryIfFailed(executeQueryExecutor, retryingTimes, sql, dataSourceIndexAndNeedRetryIfDailed);
        }
    }

    protected ResultSetWrapper executeQuery(ConnectionWrapper connectionWrapper, String sql) throws SQLException {
        StatementWrapper statement = createStatementWrapper(connectionWrapper, sql, false);
        this.resultSetWrapper = statement.executeQuery(sql);
        return this.resultSetWrapper;
    }


    public boolean getMoreResults() throws SQLException {
        return moreResults;
    }

    public int getQueryTimeout() throws SQLException {
        return queryTimeout;
    }

    public void setQueryTimeout(int queryTimeout) throws SQLException {
        this.queryTimeout = queryTimeout;
    }

    public ResultSetWrapper getResultSet() throws SQLException {
        return resultSetWrapper;
    }

    public int getUpdateCount() throws SQLException {
        return updateCount;
    }

    public Connection getConnection() throws SQLException {
        return connection;
    }

    public void cancel() throws SQLException {
        // 调用底层进行关闭
        // see com.mysql.api.StatementImpl
        this.statementWrapper.cancel();
    }

    public int getFetchSize() throws SQLException {
        return this.fetchSize;
    }

    public void setFetchSize(int fetchSize) throws SQLException {
        this.fetchSize = fetchSize;
    }

    public int getMaxRows() throws SQLException {
        return this.maxRows;
    }

    public void setMaxRows(int maxRows) throws SQLException {
        this.maxRows = maxRows;
    }

    public ResultSetWrapper getGeneratedKeys() throws SQLException {
        if (this.statementWrapper != null) {
            return this.statementWrapper.getGeneratedKeys();
        } else {
            throw new SQLException("Generated slotKeys not requested. You need to specify Statement.RETURN_GENERATED_KEYS to Statement.executeUpdate() or Session.prepareDataNodePreparedStatement().");
        }
    }

    public String getApplicationId() {
        return this.applicationId;
    }
}
